package com.central.scheduler.plugin.task.datasource.hive;

import com.central.common.constant.CommonConstant;
import com.central.common.enums.DbType;
import com.central.common.model.DatasourceInfo;
import com.central.scheduler.plugin.task.datasource.*;
import com.central.scheduler.plugin.task.util.CommonUtils;
import com.central.scheduler.spi.utils.JSONUtils;
import com.fasterxml.jackson.core.type.TypeReference;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.security.UserGroupInformation;

import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;

import static com.central.scheduler.plugin.task.datasource.PasswordUtils.decodePassword;
import static com.central.scheduler.plugin.task.datasource.PasswordUtils.encodePassword;
import static com.central.scheduler.spi.task.TaskConstants.*;

/**
 * @author Tindy
 * @date 2021/12/23
 * @describe
 */
public class HiveDatasourceProcessor extends AbstractDatasourceProcessor {


    @Override
    public ConnectionParam createConnectionParams(DatasourceInfo datasource) {
        HiveConnectionParam hiveConnectionParam = new HiveConnectionParam();
        hiveConnectionParam.setUser(datasource.getUsername());
        hiveConnectionParam.setPassword(datasource.getPassword());
        String address = datasource.getHost() + ":" + (StringUtils.isNotEmpty(datasource.getPort()) ? datasource.getPort() : "10000");
        hiveConnectionParam.setAddress(address);
        hiveConnectionParam.setDatabase(datasource.getDbName());
        hiveConnectionParam.setOther(datasource.getOther());
        hiveConnectionParam.setDriver(datasource.getDriver());
        return hiveConnectionParam;
    }

    @Override
    public String getDatasourceDriver() {
        return ORG_APACHE_HIVE_JDBC_HIVE_DRIVER;
    }

    @Override
    public String getJdbcUrl(ConnectionParam connectionParam) {
        HiveConnectionParam hiveConnectionParam = (HiveConnectionParam) connectionParam;
        StringBuilder address = new StringBuilder();
        address.append(JDBC_HIVE_2);
        address.append(hiveConnectionParam.getAddress());
        String jdbcUrl = address.toString() + "/" + hiveConnectionParam.getDatabase();
        String other = transformOther(hiveConnectionParam.getOther());
        if (StringUtils.isNotEmpty(other)) {
            jdbcUrl += ";" + other;
        }
        return jdbcUrl;
    }

    @Override
    public Connection getConnection(ConnectionParam connectionParam) throws IOException, ClassNotFoundException, SQLException {
        HiveConnectionParam hiveConnectionParam = (HiveConnectionParam) connectionParam;
        Connection connection=null;
        synchronized (CommonConstant.KRB5CONF_LOCK){
            CommonUtils.loadKerberosConf(hiveConnectionParam.getJavaSecurityKrb5Conf(),
                    hiveConnectionParam.getLoginUserKeytabUsername(), hiveConnectionParam.getLoginUserKeytabPath());
            String driver = StringUtils.isNotEmpty(hiveConnectionParam.getDriver()) ? hiveConnectionParam.getDriver() : getDatasourceDriver();
            Class.forName(driver);
            connection = DriverManager.getConnection(getJdbcUrl(connectionParam),
                    hiveConnectionParam.getUser(), decodePassword(hiveConnectionParam.getPassword()));
        }
        return connection;
    }

    @Override
    public DbType getDbType() {
        return DbType.HIVE;
    }

    private String transformOther(String other) {
        StringBuilder otherSb = new StringBuilder();
        if (!StringUtils.isBlank(other)) {
            Map map = JSONUtils.parseObject(other, new TypeReference<LinkedHashMap<String, String>>() {
            });
            if (map.size() > 0) {
                Set<String> keys = map.keySet();
                for (String key : keys) {
                    otherSb.append(String.format("%s=%s%s", key, map.get(key), ";"));
                }
                otherSb.deleteCharAt(otherSb.length() - 1);
            }
        }
        return otherSb.toString();
    }
}
